from torch.utils.data import Dataset
from pathlib import Path
from sklearn.model_selection import train_test_split
from PIL import Image
import torchvision.transforms as transforms


transformer = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
)


def loadTrainTestImgPath(base_dir, test_size=0.3):
    base_dir = Path(base_dir)
    images = []
    targets = []
    for label in base_dir.iterdir():
        target = int(label.name) - 1
        for img in label.iterdir():
            images.append(img)
            targets.append(target)
    X_train, X_test, y_train, y_test = train_test_split(
        images, targets, test_size=test_size)
    return X_train, X_test, y_train, y_test


def readImage(path):
    return Image.open(path)


class RGBImageSet(Dataset):
    def __init__(self, base_dir, test_size=0.3, inmemory=True):
        """
        base_dir: str 图片路径
        transform: 是否进行transform
        inmemory: 是否一次性将图片全部读入内存
        """
        data = loadTrainTestImgPath(base_dir, test_size)
        if inmemory:
            self.readImages(data)
        else:
            self.X_train, self.X_test, self.y_train, self.y_test = data
        self.set_train_mode = True
        self.inmemory = inmemory

    def readImages(self, data):
        X_train, X_test, self.y_train, self.y_test = data
        self.X_train = []
        for p in X_train:
            data = readImage(p)
            data = transformer(data)
            self.X_train.append(data)
        self.X_test = []
        for p in X_test:
            data = readImage(p)
            data = transformer(data)
            self.X_test.append(data)

    def trainMode(self):
        self.set_train_mode = True

    def testMode(self):
        self.set_train_mode = False

    def __getitem__(self, index):
        # 判断是训练集还是测试集
        dataList = self.X_train if self.set_train_mode else self.X_test
        targetList = self.y_train if self.set_train_mode else self.y_test
        data = dataList[index]
        target = targetList[index]
        # 如果不是inmemory模式，则data还只是path
        if not self.inmemory:
            data = readImage(data)
            # 如果需要，还应当transfor,
            data = transformer(data)
        return data, target

    def __len__(self):
        if self.set_train_mode:
            return len(self.y_train)
        else:
            return len(self.y_test)


if __name__ == "__main__":
    pass
